Critical bug fix and performance optimizations for DFlash speculative
decoding. Acceptance rate improved from 25% to 89% (matching Python
reference), throughput from 6.7 to 42 tok/s.

Root cause: hiddenNorm was declared as  without @ModuleInfo,
so its RMSNorm weight was never loaded from safetensors. The key
"hidden_norm.weight" didn't match the reflected key
"hiddenNorm.weight", leaving the weight at all-ones instead of
the trained values (~0.98). This single missing weight distorted
every draft prediction, compounding through all 5 draft layers.

Fix: Added @ModuleInfo(key: "hidden_norm") annotation, matching
the safetensors key. Also added @ModuleInfo for norm and fc for
consistency.

Performance optimizations:
- Streaming: replaced generateSync + buffered array with
  generateStreaming + Continuation, yielding tokens immediately
- Draft prefetch: launch next cycle's draft with asyncEval before
  rollback, overlapping GPU work
- Batched asyncEval: changed blocking eval() to asyncEval() for
  verify logits and hidden states
- asyncEval(committedHidden): unblocks prefetch window
- Stop token Set: precomputed O(1) lookup
- Removed double fflush, added DFlashDumper call-site guards

Submodule updates:
- mlx-swift-lm: exactSmallProjPad for quantized linear at small
  seq_len (<16), DFlash protocols, open MambaCache/ArraysCache
- mlx-swift: remove stale .air kernel files

Benchmark (Qwen3.5-27B-4bit, thinking mode, 2048 tokens):
  41.9 tok/s, 89.4% acceptance, 216 cycles

… streaming

When SSD expert streaming is active, expert weight tensors (.weight) are
replaced with zero-filled placeholders of the correct shape/dtype during
loading. Only scales and biases are loaded into RAM — the actual expert
weight data is read from SSD at runtime via pread/mmap.

RAM savings for MoE models:
  - Qwen3.6-35B-A3B: 18.4 GB → 5.1 GB (73% reduction)
  - Expert weights skipped: 16.1 GB (weight only, not scales/biases)
  - Expert scales+biases loaded: ~2 GB (needed for dequantization)

Performance on Qwen3.6-35B-A3B (512 tokens, math prompt):
  - No SSD streaming:   11.5 tok/s,  18.4 GB RAM
  - SSD streaming only: 11.5 tok/s,   5.1 GB RAM
  - SSD + DFlash:       32.2 tok/s,   5.1 GB RAM

Both streaming and non-streaming chat/text completion responses now include
a 'timings' object with:
  - predicted_per_second: generation speed in tokens/second
  - predicted_n: number of completion tokens
  - predicted_ms: total generation wall-clock time in ms

This matches llama-server's timing convention and allows clients to see
generation speed directly from the API response without external measurement.

Tests 4 configurations for Qwen3.6-35B-A3B-4bit with same math prompt:
  - Baseline (no SSD, no DFlash)
  - SSD Streaming only
  - SSD Streaming + DFlash
  - DFlash only

Results (512 tokens, 3 runs each):
  Baseline:      26.3 tok/s,  18.8 GB RAM
  SSD Streaming: 12.5 tok/s,   5.4 GB RAM
  SSD + DFlash:  33.3 tok/s,   7.4 GB RAM  ← best tradeoff
  DFlash only:  125.4 tok/s,  20.0 GB RAM

- Add StreamableMoE conformance to Qwen3NextModelInner
- Add LayerPartitionable conformance to Qwen3NextModelInner
- Add DFlashTargetModel conformance to Qwen3NextModel
  - dflashEmbedTokens, dflashLmHeadLogits, dflashForwardWithCapture
  - dflashGatedDeltaForward with tape recording for GDN rollback
- Add dflashForwardWithTape to Qwen3NextGatedDeltaNet
- Add bridge file Qwen3Next+DFlash.swift
- Short prompt works: 68.8% acceptance, 9.8 GB RAM (vs 45 GB full load)
- Longer runs crash — likely Metal watchdog on 512-expert SSD reads

…b440)

- Bumps mlx-swift-lm submodule to b440 (tag) / 63707c0:
  fix(Gemma4Text): dispatch QuantizedKVCache correctly in LLM attention
  (merges PR #29, closes #71)

- Server.swift: expose `kv_bits` as a per-request API field
  (ChatCompletionRequest.kvBits -> GenerateParameters.kvBits)
  enabling native MLX QuantizedKVCache without a server restart.

- run_benchmark.sh: add Test 9 — QuantizedKVCache regression suite
  [1/4] kv_bits=4 short  [2/4] kv_bits=8 short
  [3/4] kv_bits=4 long (KV-sharing path)  [4/4] baseline

  Test 9 passed on mlx-community/gemma-4-26b-a4b-it-4bit.

README.md:
- Added '🔧 Per-Request API Parameters' section with kv_bits table,
  kv_bits vs --turbo-kv comparison table, and curl usage example
- Clarified --turbo-kv CLI entry: 'activates after 2048 tokens, server-wide'

Server.swift:
- Added kv_bits input validation (only nil/4/8 accepted; returns 400 otherwise)
- Bypass prompt cache restore when kv_bits is set (prevents unsafe mixing of
  QuantizedKVCache and KVCacheSimple states across requests)
- Bypass prompt cache save when kv_bits is set (same safety reason)

run_benchmark.sh (Test 9):
- Corrected header comment to match actual assertions (removed false ≥20 token
  and multi-turn claims; stated actual ≥3 token / non-empty checks)
- Added explicit SERVER_READY flag + post-loop failure with log dump
- Widened thinking-block regex to handle both <|channel|>thought and <|channel>thought

- Replace 🧠 with 📡 heading emoji
- Rewrite as structured tables (Text / Vision / Audio) with all 50+ model
  families derived from the actual MLXLLM + MLXVLM model file inventory
- LLM table: Gemma, Qwen, Phi, Mistral, Llama, GLM, DeepSeek, Falcon,
  LFM2, OLMo, Granite, SmolLM3, InternLM2, Cohere, Jamba, Exaone, MiMo,
  Ernie, Baichuan, Bailing, NemotronH, Starcoder2, OpenELM, BitNet,
  MiniMax, Apertus/AfMoE, MiniCPM, Qwen3Next
- VLM table: Gemma4, Gemma3, Qwen3-VL, Qwen2-VL/2.5-VL, LFM2-VL,
  Pixtral, PaliGemma, Idefics3, Mistral3, FastVLM, SmolVLM2, GlmOcr, QwenVL
- ALM table: Gemma-4-e4b only (factually correct — Qwen2-Audio removed;
  it was never wired into the audio pipeline here)

…reverted from main in 50c3732)

fix: Gemma-4 QuantizedKVCache + kv_bits API + Test 9 (mlx-swift-lm b440)

… fix CORS/parallel test gaps

- Server.swift: add defer-based heartbeat cleanup in both handleChatStreaming and
  handleTextStreaming so heartbeatTask is always cancelled on any exit path
  (client disconnect during prefill no longer leaks the heartbeat task)
- ServerSSETests.swift: add missing import Foundation for Data/JSONSerialization
- test-server.sh Test 32: fail on empty curl response instead of false-passing
- test-server.sh Test 33: use conditional curl; fail if request fails entirely
- test-server.sh Test 34: redirect CORS preflight to CORS_PORT (--cors server)
  instead of the main server which has no CORS middleware
- test-server.sh Test 35: spin up a dedicated --parallel 2 server so concurrent
  requests actually overlap and stress the global hook under real parallelism
- test-opencode.sh: capture opencode exit code separately; classify parse errors
  vs acceptable non-zero exits to prevent false passes

…ch in Tests 32-33

The new conditional curl patterns in Tests 32 and 33 combined with the
existing set -euo pipefail caused the script to abort when grep found
no match (exit 1) in the EVENT_DATA pipeline. All grep/jq calls that
may produce no output now use || true or are wrapped in if/else to
prevent premature script exit.

[codex] make OpenAI streaming strict by default

… workflow

…experts are combined

Fixes #72: on a 16GB Mac Mini M4, adding --draft-model alongside --stream-experts
caused RAM to spike to the physical limit and trigger swap, even though the draft
model is only a 4B (~3.5GB) model.

Root causes and fixes:
1. [Bug] draftConfig.lazyLoad was never set — draft weights were eagerly paged into
   unified RAM. Fix: set draftConfig.lazyLoad = true when --stream-experts is active,
   mirroring what already happens for the main model config.

2. [Bug] Memory.cacheLimit / Memory.memoryLimit were applied after both model loads,
   so neither the main nor draft model loaded under a cache budget. Fix: apply the
   SSD memory cap immediately after ExpertStreamingConfig.shared.activate() — before
   any LLMModelFactory.loadContainer() calls — so both models respect the page-cache
   limit throughout loading.

3. [Bug] physicalBudget did not account for the draft model's resident footprint,
   leaving the cap 3–4 GB too high. Fix: profile the draft model directory before
   loading and subtract its weightMemoryGB from physicalBudget in all three affected
   strategy branches (swapAssisted, layerPartitioned, early cap). A 2 GB floor guard
   prevents the budget going negative on very constrained machines.

Expected result on 16GB M4:
- Draft model weights are mmap'd (lazy) — only accessed pages in RAM
- Both models load under the ~6GB effective page-cache budget (9.6GB - 3.5GB draft)
- No swap; total RAM stays within the SSD streaming budget

…draft model

- Extract computeSSDMemoryBudget() from inline formula so it can be unit tested
  without loading a real model or touching Memory.cacheLimit
- Wire all three budget call sites to use the extracted function (no behaviour change)
- Add SSDMemoryBudgetTests.swift with 8 tests covering:
    * Baseline 16 GB / no draft (formula correctness)
    * Issue #72 regression: 16 GB + 3.5 GB draft → budget reduced by exact footprint
    * Floor guard: deeply negative raw result clamped to 2 GB
    * Floor value: confirmed at exactly 2 GB
    * Default-arg == 0 (no silent reduction without a draft model)
    * Monotonicity: larger draft → smaller or equal budget
    * Typical fleet: 24 GB and 64 GB with 3.5 GB draft

Two correctness issues flagged in inline review:

1. GiB/GB unit mismatch — weightMemoryGB is computed as bytes/1e9 (decimal GB),
   but was multiplied back to bytes using 1_073_741_824 (GiB), causing ~7% budget
   drift. Fix: use draftProfile.weightFileSizeBytes directly (exact bytes, no
   conversion needed).

2. Repeated ModelProfiler.profile() filesystem walks — the draft model directory
   was enumerated once in the early cap block and again in each strategy branch
   (swapAssisted, layerPartitioned). Fix: compute draftFootprintBytes once before
   the streamExperts block and reuse it everywhere.

Also addresses a third Copilot comment: the early SSD cap was only applied when
modelDirectory != nil, so first-run downloads were unprotected. Now the cap is
applied whenever --stream-experts is set, even if the model isn't cached yet
(modelling via the else-if branch).

All 8 SSDMemoryBudgetTests still pass.

fix(ssd-stream): prevent RAM explosion when --draft-model + --stream-experts combined (#72)

…odel (#72 follow-up)

Reporter confirmed the original fix addressed load-time RAM, but swap still
explodes during inference: OS_RAM=20.7GB / MEM_DEMAND=40.2GB on a 16GB machine.

Root cause (inference-time):
The 200GB memoryLimit sentinel is necessary for SSD streaming alone — it bypasses
MLX eval_impl's spin-wait loop when expert pages are evicted mid-graph.  However,
with speculative decoding the draft model (4B / 3GB) and main model (35B / 20GB)
alternate forward passes in tight succession.  Both models' expert pages are
demanded within the same inference cycle, combined demand ~23GB >> 16GB physical.
The 200GB sentinel provides zero back-pressure, so macOS swaps aggressively
(10+ GB observed in Activity Monitor).

Fix:
When --stream-experts + --draft-model are both set AND combinedFootprint > 70%
of physical RAM, lower memoryLimit from 200GB to physicalRAM × 1.1.  This forces
MLX to hit its hard limit sooner and evict stale expert pages more aggressively
rather than extending into swap.  A clear startup warning is also printed:

  ⚠️  SSD + draft-model RAM pressure warning:
     Main model: 20.4GB  Draft: 3.0GB  Combined: 23.4GB  Physical RAM: 16.0GB
     Speculative decoding alternates both models' forward passes.
     On this machine the combined weight exceeds physical RAM,
     causing page-cache thrashing and swap during inference.
     → Recommendation: remove --draft-model on this machine,
       or use a smaller draft model whose weights fit in
       remaining RAM after the main model's page budget (6GB).
     Memory limit set to 17GB (tight cap for MLX eviction pressure)

When combined footprint fits in RAM (e.g. smaller draft on a 32GB machine),
the 200GB sentinel is still used as before — no regression for capable hardware.

…el cache

Replace if-branch masking with metal::select for zero warp-divergence state
updates. Reorganize KernelCache from 8 flat named vars to tapeReplay[vec][msk]
and gatedDeltaTape[vec][msk] 2D arrays. Simplify dispatch call sites to
one-liner index lookups. Minor whitespace cleanup in DFlashIntermediateDumper.

… property

Add MambaSnapshotCache: lightweight O(1) snapshot-based rollback (lazy
reference capture, no GPU copy) as an alternative to RecurrentRollbackCache's
innovation-tape replay. Add dflashUseTapeRollback Bool to DFlashTargetModel
(default true) so models can opt in to either strategy. Update makeTargetCache
and arm/rollback helpers with clearer comments.

Also switch RecurrentRollbackCache.armRollback to lazy reference capture
(removes unnecessary MLX.contiguous copies on arm path).

Add DFlashKernelBench executable for isolated kernel timing. Exclude
DFlashKernelsOptimized.swift from the DFlash library target (work-in-progress
alternative kernel implementations kept for reference).

…next.sh

bench_35b.sh: save per-run raw response JSON, extract structured results into
bench_results.json (tok/s, RAM, timing per config) for downstream tooling.
Use slug variable consistently for log file naming.

Add bench_coder_next.sh for benchmarking Qwen3-Coder-Next model variants.

Move comparison tests from tests/DFlashComparison/ to tests/DFlash/, adding
DFlashBenchmark.swift, DFlashProfiler.swift, updated cosine similarity
comparison tools, and a README. Update .gitignore intermediates path.

…-draft-model (#72)

Git history audit (mlx-swift-lm):
  e6ba580 - 8.5x speedup (0.58→4.95 tok/s) from cross-projection batching (Eric Lake, M1 Ultra)
  2c71c6c - ssd-opt-v2: +4% more via persistent expert buffers (asyncEval warm path)
  2b1c653 - PAPPS N+1 prefetch permanently disabled (hurt Apple-native TPS)

README (line 245) explicitly states:
  'Speculative decoding is counterproductive for SSD-streaming MoE specifically.
   The verify pass sends N+1 tokens, each routing to *different* experts — SSD I/O
   scales with the *union* of all positions' expert selections.'

Strategy (not a hard error):
When --stream-experts + --draft-model are combined:
  - Auto-cap --num-draft-tokens to 1 (verify pass = 2 positions, not N+1)
  - At 1 draft token: fan-out is 2× SSD I/O (vs 5× at default 4 tokens)
  - If acceptance rate ≥ 50% (typical for same-family models), net TPS is positive
  - Print a clear advisory so users understand the tradeoff
  - Persistent expert buffers (~5 GB warm path, ssd-opt-v2) are PRESERVED —
    no regression to Eric Lake's M1 Ultra benchmark

What is NOT changed:
  - SwitchLayers.swift warm path: untouched (idx.size <= 32 guard intact)
  - ExpertStreamingConfig: no new flags added (reverted failed hasDraftModel attempt)
  - computeSSDMemoryBudget() + cacheLimit logic from load-time fix: intact
  - Tight memoryLimit sentinel (physicalRAM × 1.1) when combined > 70% RAM: intact

Test coverage (18 tests, 0 failures):
  SSDDraftStrategyTests (10 new):
    - Fan-out arithmetic: 4 draft tokens → 5× I/O, 1 token → 2× I/O
    - Auto-cap fires only when streamExperts + draftModel + numDraftTokens > 1
    - Auto-cap does NOT fire for solo SSD streaming or pure RAM speculative decoding
    - Net throughput model: 70% acceptance at 2× fan-out is net positive
    - memoryLimit sentinel selection: tight cap on 16 GB, sentinel on 64 GB
  SSDMemoryBudgetTests (8 existing): all pass, no regressions

…sion

Three-check E2E test for the --stream-experts + --draft-model fix:

  [1/3] Auto-cap guard: verifies server log contains the 'auto-capping'
        warning, proving numDraftTokens was reduced from 4 to 1 at startup

  [2/3] RAM guard: measures vm_stat peak RAM during inference and fails
        if it exceeds 80% of physical RAM (the indicator that exposed the
        original swap explosion on reporter's 16GB M4 Mini)

  [3/3] Inference: verifies the combination still produces valid content
        (not crashed/empty), proving functional correctness

Uses small models (Qwen3.5-4B main + Qwen3.5-0.8B draft) — same
parameter-class proportions as the reporter's 35B+4B scenario but
runnable on any machine without 35B weights.

Run: ./run_benchmark.sh → option 10

Replacing the stdin pipe approach with an env var so child invocations
from Test 0's automated matrix loop skip the interactive menu entirely.
The previous echo-pipe was consumed by the 'read suite_opt' prompt but
any subsequent reads (model selection) had no input, causing the script
to fall through to option 3 by default.

When SUITE_OPT is set (automated matrix mode), skip all menu echoes
and the read prompt entirely. Child processes now run silently with
only test-relevant output.

Both test-speculative.sh and test-dflash.sh grep for 'Using speculative
decoding' in the server log to confirm the speculative path was activated.
This string was never emitted — the tests were checking a log line that
didn't exist, causing speculative-decoding and dflash-speculative-decoding
CI jobs to always fail on Test 1.

Fix: emit the exact expected log line:
  - Standard spec: after draft model is loaded successfully
  - DFlash spec: at generation dispatch in Server.swift

Server log now contains all strings the tests grep for:
  ✅ 'Draft model loaded successfully'
  ✅ 'Using speculative decoding'
  ✅ 'speculative decoding' (for test-speculative-eval.sh)

test-dflash.sh grepped for:
  1. 'Draft model loaded successfully' — only emitted by standard draft path,
     not DFlash path which has its own 'DFlash draft model loaded' message
  2. 'Using speculative decoding' — not emitted by DFlash path at all
  3. 'speculative decoding' — was present but test was failing on (1)

Add both required lines immediately after DFlash draft model weights load,
mirroring the standard speculative decoding path. The streaming failures
('missing [DONE] sentinel') were downstream of the model-not-found state
caused by the load log mismatch, not an inference bug.

Adds Sources/SwiftLM/{Qwen3,Qwen3MoE,Llama}+DFlash.swift — each
declares the DFlashTargetModel protocol conformance and delegates to
the model's public callCapturing / embedTokens / lmHead
(now on *ModelInner via mlx-swift-lm b453).

Coverage:
  Qwen3Model      → Qwen3-8B and similar dense Qwen3 variants
  Qwen3MoEModel   → Qwen3-Coder-30B-A3B and other Qwen3 MoE variants
  LlamaModel      → Meta-Llama-3.x, Mistral, and Llama-family models
  Qwen35MoEModel  → already covered via Qwen35Model inheritance
  Qwen36MoE       → no separate Swift class found; uses Qwen35MoE path

Co-authored-by: clandestine.eth <96172957+0xClandestine@users.noreply.github.com>

Gemma4 omni (5.2GB) on a 7.5GB runner is tight. After other CI jobs
have run and filled the model cache, available RAM can drop below the
threshold needed for stable Metal command buffer execution, causing
sporadic GPU timeout crashes (kIOGPUCommandBufferCallbackErrorTimeout).

Add a vm_stat-based preflight check: if available+inactive RAM < 2.5GB,
exit 0 (skip) instead of crashing the whole run.

This reverts commit 9fc993c.

Own DeepSeek V3 (deepseek_v3 / kimi_k25) and Kimi Linear (kimi_linear)
model implementations directly in SwiftLM so DFlashTargetModel conformance
is available without any upstream submodule changes.

- DeepseekV3DFlash.swift: full DSV3Config + model with callCapturing
- KimiLinearDFlash.swift: hybrid KDA/MLA Kimi 2.6 model with DFlash
- DFlashModelRegistry.swift: registers all three model types via
  LLMTypeRegistry.shared.registerModelType() at startup
- Server.swift: call registerDFlashModelTypes() before model loading

…in-loops

Use @ModuleInfo(key: "model") on the inner model property so weights
at model.* paths are found correctly. Also use @ModuleInfo(key: "norm")
for norm layers initialized in init() so their weights are tracked.

… limit

DeepseekV3DFlash.sanitize():
- Strip 'language_model.' wrapper prefix present in kimi_k25 and some
  other HuggingFace exports so weight keys resolve to model.* paths
- After stacking per-expert weights into switch_mlp, remove the original
  experts.N.* keys to prevent verify: .noUnusedKeys crash
- Generalize layer filter to use numHiddenLayers instead of hardcoded 61

Server.run():
- Raise RLIMIT_NOFILE to 4096 at startup; large sharded models (kimi_k25
  has 182 safetensor shards) exhaust the default macOS limit of 256

…prevent GPU timeouts

- Move MLX_MAX_OPS_PER_BUFFER=50 to top of run() before Metal init
- Enable --stream-experts automatically on <12GB machines in test-dflash.sh
  so weights are paged via mmap/pread instead of macOS VM swap
- Auto-cap draft tokens to 1 under SSD streaming (minimal fan-out)
- Always compute draftFootprintBytes regardless of --stream-experts flag

* feat: bump mlx-swift-lm submodule for DeepSeek-V4 support

Points mlx-swift-lm to feat/deepseek-v4 branch (SharpAI/mlx-swift-lm#33)
which adds DeepseekV4.swift and registers the deepseek_v4 model type.

* feat: DeepSeek-V4-Flash benchmark results + profiler improvements

- README: add DeepSeek-V4-Flash (126GB Q3) benchmark table for M5 Pro 64GB
  SSD+TurboQuant delivers 4.16 tok/s at 40K context (13x vs plain SSD Stream)
- profile_runner.py: track peak GPU InUse via background polling thread (0.5s)
  instead of single post-generation snapshot; rename gpu_in_use → gpu_in_use_peak
  throughout; add separate GPU_InUse peak visualization section
- run_benchmark.sh: add Thump604/DeepSeek-V4-Flash-MLX-Q3-mixed-gs128-affine
  to Test 1 model list (option 11)
- mlx-swift-lm: bump submodule to 8a8da29 (attn_sink dtype fix)

* chore: bump mlx-swift-lm submodule to b463 (DeepSeek-V4 merged to main)

…on (#81)

…83)

feat: add DFlash speculative decoding

Merges ericjlake's prompt-cache fixes from PR #85, resolving conflicts
with the DFlash integration (PR #78).

Changes from ericjlake:
- MambaCache safety gate + KVCacheSimple T-dim slice in save()
- ndim >= 3 guard in minCachedSeqLen scan
- Spec-decode short-circuit ordering (check before cache restore)
- README: Qwen3-A3B full-RAM perf table (M1 Ultra 64 GB)

Conflict resolution:
- README.md: kept both Qwen3-A3B and DeepSeek-V4 perf tables
- Server.swift save(): kept existing MambaCache early return + new T-dim slice
- Server.swift decision branch: combined spec-decode-first + skipPromptCache (kvBits)

Closes #84.
Co-authored-by: Eric Lake <ericjlake@users.noreply.github.com>